# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: inkyu

import hashlib
import hmac
import time

from urllib.parse import urlparse
from requests.auth import AuthBase


class APIKeyAuthWithExpires(AuthBase):
  def __init__(self, api_key: str, api_secret: str):
    self._api_key = api_key
    self._api_secret = api_secret

  def generate_signature(self, method: str, url: str, nonce: int, data=''):
    # Remove the base
    parsed_url = urlparse(url)
    path = parsed_url.path
    if parsed_url.query:
      path = path + '?' + parsed_url.query
    if isinstance(data, (bytes, bytearray)):
      data = data.decode('utf8')
    message = method + path + str(nonce) + data
    signature = hmac.new(bytes(self._api_secret, 'utf8'),
                         bytes(message, 'utf8'),
                         digestmod=hashlib.sha256).hexdigest()
    return signature

  def __call__(self, r):
    expires = int(round(time.time()) + 5)
    r.headers['api-expires'] = str(expires)
    r.headers['api-key'] = self._api_key
    r.headers['api-signature'] = self.generate_signature(r.method, r.url, expires, r.body or '')
    return r
